import cv2
import numpy as np
import os
import csv
import math
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt
import concurrent.futures
from functools import partial

# --- 全局变量用于鼠标回调 ---
line_points = []
calibration_image_display = None

# --- Matplotlib 中文和负号显示配置 ---
# 确保您的系统安装了支持中文的字体，例如 'SimHei', 'Microsoft YaHei' 等
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False   # 解决保存图像是负号'-'显示为方块的问题


def mouse_callback(event, x, y, flags, param):
    """处理鼠标点击事件以进行标定"""
    global line_points, calibration_image_display
    
    if event == cv2.EVENT_LBUTTONDOWN:
        if len(line_points) < 2:
            line_points.append((x, y))
            cv2.circle(calibration_image_display, (x, y), 5, (0, 255, 0), -1)
            if len(line_points) == 2:
                cv2.line(calibration_image_display, line_points[0], line_points[1], (0, 255, 0), 2)
            cv2.imshow("标定窗口 - 请画一条线", calibration_image_display)

def get_pixel_to_mm_ratio(image_path):
    """通过用户交互获取像素到毫米的转换比例。"""
    global line_points, calibration_image_display
    line_points = []

    image = cv2.imread(image_path)
    if image is None:
        print(f"错误：无法读取用于标定的图像 {image_path}")
        return None

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    calibration_image_display = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)

    cv2.namedWindow("标定窗口 - 请画一条线")
    cv2.setMouseCallback("标定窗口 - 请画一条线", mouse_callback)

    print("\n--- 标定步骤 ---")
    print("1. 在弹出的 '标定窗口' 中，用鼠标左键点击两次，画一条你知道其实际长度的直线。")
    print("2. 画线完成后，请不要关闭窗口。")
    print("3. 回到这个命令行窗口，输入该直线的实际长度。")

    while len(line_points) < 2:
        if cv2.waitKey(20) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            print("标定被用户取消。")
            return None
        cv2.imshow("标定窗口 - 请画一条线", calibration_image_display)

    try:
        real_length_mm = float(input("请输入您所画直线的实际长度 (单位: 毫米)，然后按 Enter: "))
    except ValueError:
        print("输入无效，请输入一个数字。")
        cv2.destroyAllWindows()
        return None

    cv2.destroyAllWindows()

    p1, p2 = line_points
    pixel_distance = math.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)

    if real_length_mm <= 0:
        print("实际长度必须大于0。")
        return None

    pixels_per_mm = pixel_distance / real_length_mm
    print(f"标定完成: 像素距离 = {pixel_distance:.2f} px, 实际长度 = {real_length_mm} mm")
    print(f"计算出的比例尺: {pixels_per_mm:.2f} 像素/毫米")
    
    return pixels_per_mm

def analyze_and_plot_vibration(csv_path, pixels_per_mm, frame_rate, output_folder):
    """读取质心数据，进行振动分析并绘图。"""
    # 修改：从CSV读取X坐标数据
    x_coords_pixel = []
    with open(csv_path, 'r', encoding='utf-8') as csvfile: 
        reader = csv.reader(csvfile)
        next(reader)  # 跳过表头
        for row in reader:
            x_coords_pixel.append(float(row[1])) # 读取X坐标 (索引为1)

    if not x_coords_pixel:
        print("未能从CSV文件中读取到数据。")
        return

    x_coords_pixel = np.array(x_coords_pixel)
    mean_x_pixel = np.mean(x_coords_pixel)
    displacement_mm = (x_coords_pixel - mean_x_pixel) / pixels_per_mm

    N = len(displacement_mm)
    T = 1.0 / frame_rate
    time = np.linspace(0.0, N * T, N, endpoint=False)

    plt.figure(figsize=(12, 6))
    plt.plot(time, displacement_mm)
    plt.title('X方向振动位移-时间图') # 更新标题
    plt.xlabel('时间 (s)')
    plt.ylabel('位移 (mm, X方向)') # 更新Y轴标签
    plt.grid(True)
    plt.savefig(os.path.join(output_folder, 'displacement_vs_time_X.png')) # 更新文件名
    
    displacement_no_dc = displacement_mm - np.mean(displacement_mm)
    yf = fft(displacement_no_dc)
    xf = fftfreq(N, T)[:N//2]
    amplitude = 2.0/N * np.abs(yf[0:N//2])

    plt.figure(figsize=(12, 6))
    plt.plot(xf, amplitude)
    plt.title('X方向频域分析 (FFT)') # 更新标题
    plt.xlabel('频率 (Hz)')
    plt.ylabel('振幅 (mm, X方向)') # 更新Y轴标签
    plt.grid(True)
    
    if len(amplitude) > 1:
        peak_index = np.argmax(amplitude[1:]) + 1 
        peak_frequency = xf[peak_index]
        peak_amplitude = amplitude[peak_index]
        plt.axvline(peak_frequency, color='r', linestyle='--', label=f'主频率: {peak_frequency:.2f} Hz')
        plt.legend()
        print("\n--- 振动分析结果 ---")
        print(f"主要振动频率为: {peak_frequency:.2f} Hz")
        print(f"对应的振幅为: {peak_amplitude:.4f} mm")
    else:
        print("\n--- 振动分析结果 ---")
        print("数据点不足，无法分析频率。")

    plt.savefig(os.path.join(output_folder, 'frequency_analysis_X.png')) # 更新文件名
    print(f"分析图表已保存到: {output_folder}")

def process_single_image(file_info, input_folder, output_folder):
    """
    处理单个图像文件的函数，设计为在线程池中运行。
    """
    i, filename, should_save = file_info
    file_path = os.path.join(input_folder, filename)
    image = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

    if image is None:
        return None

    if len(image.shape) == 2:
        gray_image = image
        image_for_drawing = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) if should_save else None
    elif len(image.shape) == 3:
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        image_for_drawing = image if should_save else None
    else:
        return None
        
    _, binary_image = cv2.threshold(gray_image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    kernel = np.ones((5, 5), np.uint8)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel, iterations=2)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_CLOSE, kernel, iterations=3)
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        M = cv2.moments(largest_contour)
        cX, cY = (0, 0)
        if M["m00"] != 0:
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
        else:
            x, y, w, h = cv2.boundingRect(largest_contour)
            cX, cY = x + w // 2, y + h // 2

        if should_save and image_for_drawing is not None:
            cv2.circle(image_for_drawing, (cX, cY), 7, (0, 0, 255), -1)
            cv2.putText(image_for_drawing, f"({cX}, {cY})", (cX - 40, cY - 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            output_path = os.path.join(output_folder, filename)
            cv2.imwrite(output_path, image_for_drawing)
        
        return (filename, cX, cY)
    return None

def process_images_in_folder(input_folder, output_folder, num_images_to_save):
    """使用线程池加速处理文件夹中的所有图像。"""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        print(f"创建输出文件夹: {output_folder}")

    try:
        file_list = [f for f in sorted(os.listdir(input_folder)) if f.lower().endswith(".bmp")]
        if not file_list:
            print(f"错误：在 '{input_folder}' 中未找到 .bmp 文件。")
            return None
    except FileNotFoundError:
        print(f"错误：找不到输入文件夹 '{input_folder}'。请检查路径是否正确。")
        return None

    total_files = len(file_list)
    save_indices = set(np.linspace(0, total_files - 1, num_images_to_save, dtype=int))
    
    tasks = [(i, filename, i in save_indices) for i, filename in enumerate(file_list)]
    
    print(f"开始使用多线程处理 {total_files} 个文件...")
    
    # 使用 partial 来固定 process_single_image 的部分参数
    worker_func = partial(process_single_image, input_folder=input_folder, output_folder=output_folder)
    
    results = []
    # 使用ThreadPoolExecutor来并行处理任务
    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
        try:
            from tqdm import tqdm
            # 使用tqdm显示进度条
            results = list(tqdm(executor.map(worker_func, tasks), total=total_files, desc="图像处理中"))
        except ImportError:
            print("提示: 未找到 'tqdm' 库。如需显示进度条，请运行: pip install tqdm")
            results = list(executor.map(worker_func, tasks))

    # 过滤掉处理失败的结果 (返回值为 None)
    centroids_data = [res for res in results if res is not None]
    
    # executor.map 会保持任务的顺序，所以结果已经是按文件名排序的
    print(f"处理完成。共成功分析了 {len(centroids_data)} 张图像。")
    
    csv_path = os.path.join(output_folder, 'centroids.csv')
    try:
        with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['文件名', '质心_X', '质心_Y'])
            writer.writerows(centroids_data)
        print(f"质心数据已成功保存到: {csv_path}")
        return csv_path
    except IOError as e:
        print(f"错误：无法写入CSV文件 {csv_path}。原因: {e}")
        return None

# --- 主程序 ---
if __name__ == "__main__":
    input_directory = "F:/use/500s5ms/C001H001S0001"
    FRAME_RATE = 500
    NUM_IMAGES_TO_SAVE = 10

    parent_dir = os.path.dirname(os.path.abspath(input_directory))
    output_directory = os.path.join(parent_dir, "processed_output")

    first_image_path = None
    try:
        for f in sorted(os.listdir(input_directory)):
            if f.lower().endswith('.bmp'):
                first_image_path = os.path.join(input_directory, f)
                break
    except FileNotFoundError:
        print(f"错误：找不到输入文件夹 '{input_directory}'。")
        exit()

    if not first_image_path:
        print(f"错误：在文件夹 '{input_directory}' 中没有找到任何 .bmp 图像。")
        exit()
    
    pixels_per_mm = get_pixel_to_mm_ratio(first_image_path)

    if pixels_per_mm:
        centroids_csv_file = process_images_in_folder(input_directory, output_directory, NUM_IMAGES_TO_SAVE)
        if centroids_csv_file:
            analyze_and_plot_vibration(centroids_csv_file, pixels_per_mm, FRAME_RATE, output_directory)
            print("\n所有处理和分析已完成！")
